#!/usr/bin/python

# Imports
import sys,  os
import pygame
import opencv
from opencv.cv import *
from opencv.highgui import *
from data import *
import scene
from inputbox import Inputbox
from conf import countdown_seconds, pictures_taken, max_amount_of_pics, amount_of_matches

class FaceDetect(object):
    """Used to detect, capture, save, update and match faces.
    Can also draw a live stream of the (web)cam."""

    def __init__(self, screen, sql, scene = None):
        """Initialization of FaceDetect.
        Creates the necessary global parameters and creates 2 directories when not existing.
        Parameters: display screen - drawing screen
                    SQL sql - instance of SQL()
                    Scene scene - instance of Scene()"""
        print 'Initializing FaceDetect'
        self.storage = cvCreateMemStorage(0)
        self.capture = cvCreateCameraCapture(0)
        
        face_cascade_name =  "haarcascade_frontalface_default.xml"
        face_cascade_name = os.path.join(DATADIR, face_cascade_name)
        self.face_cascade = cvLoadHaarClassifierCascade(face_cascade_name, cvSize(0, 0))
        
        self.cam_width = 320
        self.cam_height = 240
    
        cvSetCaptureProperty(self.capture, CV_CAP_PROP_FRAME_WIDTH, self.cam_width)
        cvSetCaptureProperty (self.capture, CV_CAP_PROP_FRAME_HEIGHT, self.cam_height)
        
        self.modulo = 10
        self.seconds = countdown_seconds
        self.may_match = False
        self.keepGoing = True
        self.scene = scene
        self.users = []
        self.captures = []
        self.user_id = None
        
        self.img_x_scale = 1.0
        self.img_y_scale = 1.3
        self.screen = screen
        self.sql = sql
        
        if not os.path.isdir(BASEDIR):
            os.mkdir(BASEDIR)
        if not os.path.isdir(TMPDIR):
            os.mkdir(TMPDIR)
  
    def equalizeHist(self, img):
        """Shrinks, normalizes the brightness and increases contrast of the image.
        Returns an IplImage.
        Parameters: IplImage img - image of frame of the (web)cam"""
        gray = cvCreateImage(cvSize(img.width,img.height), 8, 1)
        small_img = cvCreateImage(cvSize(cvRound (img.width/self.img_x_scale),
                                           cvRound(img.height/self.img_y_scale)), 8, 1)
        cvCvtColor(img, gray, CV_RGB2GRAY)
        cvResize(gray, small_img)
        cvEqualizeHist(small_img, small_img)
        cvClearMemStorage(self.storage)
  
        return small_img
    
    def detect_face(self, img):
        """Tries to detect a face and if detected drawing an Ellipse around it.
        Parameters: IplImage img - image of frame of the (web)cam"""
        small_img = self.equalizeHist(img)
        faces = cvHaarDetectObjects(small_img, self.face_cascade, self.storage, 
                                    1.1, 3, 0, cvSize(50,  50))
        
        if faces:
            for r in faces:
                x, y = r.x - 15, r.y - 15
                width, height = int(r.x + (r.width * self.img_x_scale) + 10), int(r.y + (r.height * self.img_y_scale) + 30)
                pt1, pt2 = cvPoint(x,  y), cvPoint(width,  height)
                cvEllipse(img, cvPoint(x + ((width - x ) / 2), y + ((height - y) / 2)), cvSize(((width - x ) / 2), ((height - y) / 2)), 0, 0, 360, CV_RGB(0,0,255), 1, 8, 0)
                rect = cvRect(x, y, width - x, height - y)
                return rect, img
        return None, img
    
    def get_list_max(self, index_list, count_list, matches_list):
        """Returns a list of users with the most matches. 
        Can only be used together with get_best_matches().
        Parameters: list index_list - list of index numbers refering to the same index in matches_list
                    list count_list - list of count numbers how often each user occurs.
                    list matches_list - list of users (each matching user occurs once in this list)
        The indexes of all lists correspond which each other. So index_list[0] = count_list[0] = match_list[0]"""
        index = 0
        if len(matches_list) > 1:
            for cl in count_list:
                if cl == max(count_list):
                    index_list.append(matches_list[index])
                    count_list.remove(count_list[index])
                    matches_list.remove(matches_list[index])
                    break
                index += 1
        elif len(matches_list) == 1:
            index_list.append(matches_list[count_list.index(max(count_list))])
            count_list.remove(count_list[index])
            matches_list.remove(matches_list[index])
        
        return index_list, count_list, matches_list
    
    def get_best_matches(self, matches):
        """Returns the best matches based on how often the match is found in the list.
        The amount returned can be specified in conf.py.
        Parameters: list matches - list of all matching users, meaning that an user can occur more then once."""
        count = 1
        count_list = []
        matches_list = []
        index_list = []
        
        for match in matches:
            if match not in matches_list:
                count_list.append(matches.count(match))
                matches_list.append(match)
        
        if len(matches) > 0:
            while len(index_list) < amount_of_matches:
                if len(matches_list) == 0:
                    break
                else:
                    index_list, count_list, matches_list = self.get_list_max(index_list, count_list, matches_list)
        
        return index_list
    
    def get_most_recent_faces(self):
        """Returns all face_ids of the the three most recent pictures of each user."""
        users = self.sql.get_all_users()
        face_ids = []
        for user in users:
            user_id = user[0]
            tmp_faces = self.sql.get_most_recent_faces(user_id)
            for user_face in tmp_faces:
                face_ids.append(user_face)
                    
        return face_ids
    
    def matchTemplate(self, capture, src):
        """Compares a template against overlapped image regions.
        Parameters: IplImage capture - capture image saved in '/tmp'.
                    IplImage src - source image which capture is compared to."""
        # Setting image properties
        temp_width  = src.width - capture.width + 1
        temp_height = src.height - capture.height + 1
        
        # Create image for template matching
        temp_diff = cvCreateImage(cvSize(temp_width, temp_height), IPL_DEPTH_32F, 1)   
        
        # Matching template
        cvMatchTemplate(src, capture, temp_diff, CV_TM_CCORR)
      
        # Getting value out of temp_diff
        temp_diff = temp_diff[0]
        
        return temp_diff
    
    def matchShapes(self, capture, src):
        """Compares two shapes.
        Parameters: IplImage capture - capture image saved in '/tmp'.
                    IplImage src - source image which capture is compared to."""
        # Creating images for grayscale
        src_gray = cvCreateImage(cvSize(src.width, src.height), src.depth,1)
        cap_gray = cvCreateImage(cvSize(capture.width, capture.height), capture.depth, 1)
        
        # Making images gray
        cvCvtColor(src, src_gray, CV_RGB2GRAY)
        cvCvtColor(capture, cap_gray, CV_RGB2GRAY)
        
        # Matching shapes
        shape_diff = cvMatchShapes(src_gray, cap_gray,  1)
        
        return shape_diff
        
    def match_all(self, face_ids):
        """Going either to the 'User List' window (when no matches were found) 
        or to the 'Match' window showing a variable amount of the found matches on the screen.
        Matches all capture(s) with the three most recent pictures of each user.
        Parameters: list face_ids - user faces that are going to be matches"""
        allmatches = []
        shape_results = []
        temp_results = []
        
        print 'Matching captures with most recent pictures of each user.'
        for capture in self.captures:
            shape_diff, temp_diff = self.match(capture)
            print 'Shape_diff:', shape_diff
            print 'Temp_diff:', temp_diff
                
            if shape_diff != None:
                shape_results = self.sql.get_closest_shape_diffs(shape_diff, face_ids)
                for user_id in shape_results:
                    allmatches.append(self.sql.get_user(user_id = user_id))
            if temp_diff != None:
                temp_results = self.sql.get_closest_temp_diffs(temp_diff, face_ids)
                for user_id in temp_results:
                    allmatches.append(self.sql.get_user(user_id = user_id))
        
        list = self.get_best_matches(allmatches)
        
        if allmatches == []:
            self.scene.user_list_window(False)
        else:
            self.scene.match_window(list)
        
        for user in allmatches:
            print 'Matching user:', user
    
    def match(self, img, src = 'fingerprint.png'):
        """Matches the image with the source image.
        Returns the values that were created by matching.
        Parameters: string img - name of image saved in '/tmp'.
                    string src - name of source image"""
        shape_diff = None
        temp_diff = None
        try:
            print 'Matching %s with %s' % (src, img)
            capture = cvLoadImage(img)
            src = cvLoadImage(os.path.join(DATADIR, src))
            shape_diff = self.matchShapes(capture, src)
            temp_diff = self.matchTemplate(capture, src)
        except:
            pass
        
        return shape_diff, temp_diff

    def save_tmp_image(self, img):
        """Saves the capture(s), taken during the login sequence, in '/tmp' so they can be used later on.
        Parameters: IplImage img - image of frame of the (web)cam"""
        try:
            face, img = self.detect_face(img)
            if face:
                name = 1
                self.filename = os.path.join(TMPDIR, str(name) + ".png")
                while os.path.exists(self.filename):
                    name += 1
                    self.filename = os.path.join(TMPDIR, str(name) + ".png")
                ROI = cvGetSubRect(img, face)
                if self.detect_face(ROI):
                    dst = cvCreateImage(cvSize(95,  135), ROI.depth, ROI.nChannels)
                    cvResize(ROI, dst)
                    
                    width = dst.width / 2
                    height = dst.height / 2
                    print width, height
                    while width < (dst.width / 2 + 25) and height < (dst.height / 2 + 25):
                        cvEllipse(dst, cvPoint(dst.width / 2, dst.height / 2), cvSize(width, height), 0, 0, 360, CV_RGB(0,255,0), 2, 8, 0)
                        width += 1
                        height += 1
                    
                    cvSaveImage(self.filename,  dst)
                    self.captures.append(self.filename)
                    print "Saving temp_pic: ", self.filename
                else:
                    print "No face detected"
            else:
                print "No face detected"
        except Exception,  info:
            print 'Error - Saving temp_pic:', info

    def save_diffs(self, tmp_filename, pic_name = None):
        """Saves the diff values in users_faces and the picture in 'capture/'
        The maximum rows allowed is specified in conf.py, when reached the oldest will get replaced by a newer.
        Parameters: string tmp_filename - name of image saved in '/tmp'.
                    string pic_name - if None then the pic_name is created else it is used for updating."""
        user_id = self.user_id
        shape_diff, temp_diff = self.match(tmp_filename)
        img = cvLoadImage(tmp_filename)
        if not pic_name:
            pic_name = str(user_id) +  "_" + str(self.sql.count_users_faces(user_id) + 1) +  ".png"
        filename = os.path.join(BASEDIR, pic_name)
        if self.sql.count_users_faces(user_id) == max_amount_of_pics:
            print 'Maximum amount (%s) of pictures reached.' % max_amount_of_pics
            oldest_face_id = self.sql.get_oldest_face_id(user_id)
            old_pic_name = self.sql.get_pic_name(oldest_face_id)
            filename = os.path.join(BASEDIR, old_pic_name)
            self.sql.update_user_face(shape_diff, temp_diff, old_pic_name, oldest_face_id)
        else:
            self.sql.create_user_face(shape_diff, temp_diff, pic_name, user_id)
    
        cvSaveImage(filename, img)
        print "Saved picture path:", filename
    
    def count_down(self):
        """Counts from the number specified in conf.py till 0. 
        The numbers are shown on the screen when used."""
        font = pygame.font.Font(None, 100)
        if not self.draw_seconds == 0:
            self.count += 1
            ren = font.render(str(self.draw_seconds), 1, (255, 255, 255))
            self.screen.blit(ren, (self.screen.get_width() / 2 - ren.get_width() / 2, self.screen.get_height() / 2 - ren.get_height() / 2))
            if (self.count % self.modulo) == 0:
                    self.draw_seconds -= 1
    
    def reset_countdown(self):
        """Resets the countdown to it's original number of seconds."""
        self.count = 1
        self.draw_seconds = self.seconds
    
    def cancel_countdown(self):
        """Cancels the countdown by setting the amount of seconds left to 0."""
        self.count = None
        self.draw_seconds = 0

    def take_pictures(self, frame):
        """Works together with the countdown and takes a variable amount of pictures after the countdown is complete.
        The amount of pictures taken can be modified in conf.py.
        Parameters: IplImage frame - frame of the (web)cam"""
        if self.count:
            if self.count >= (self.modulo * self.seconds) and self.count <= (self.modulo * (self.seconds + pictures_taken)):
                if self.count >= (self.modulo * (self.seconds + 1)):
                    if (self.count % self.modulo) == 0:
                        self.save_tmp_image(frame)
                    if (self.count % (self.modulo * (self.seconds + pictures_taken))) == 0:
                        self.may_match = True
                self.count += 1
            elif self.may_match:
                self.may_match = False
                face_ids = self.get_most_recent_faces()
                self.match_all(face_ids)
    
    def set_user_id(self, user_id):
        """Sets the global user_id parameter.
        Parameters: int user_id - user identification code"""
        print 'user_id set to', user_id
        self.user_id = user_id
    
    def draw_cam(self):
        """Draws a frame of the (web)cam on the screen.
        When used in a loop it looks like a live stream.
        Returns a modified frame of the (web)cam."""
        small_frame = self.get_cam_frame()
        
        face, img = self.detect_face(small_frame)
        
        surf = convert_image(img)
        self.screen.blit(surf, ((self.screen.get_width() / 2) - (img.width / 2), (self.screen.get_height() / 3) - (img.height / 2)))
        
        return small_frame
    
    def login(self):
        """Returns the frame of the (web)cam while counting down and taking picture(s) afterwards. """
        small_frame = self.get_cam_frame()
        
        self.count_down()
        self.take_pictures(small_frame)
        
        return small_frame
    
    def get_cam_frame(self):
        """Creates and returns the actual (web)cam frame."""
        if not self.face_cascade:
            print "Error - Could not load classifier cascade"
            sys.exit(-1)
        
        frame = cvQueryFrame(self.capture)
        small_frame = cvCreateImage(cvSize(self.cam_width,  self.cam_height), frame.depth, frame.nChannels)
        cvResize(frame, small_frame)
        cvFlip(small_frame, small_frame, 1)
        
        return small_frame
